import OpenAI from "openai";
import { Client } from "@modelcontextprotocol/sdk/client/index.js";
import { StdioClientTransport } from "@modelcontextprotocol/sdk/client/stdio.js";
import { StreamableHTTPClientTransport } from "@modelcontextprotocol/sdk/client/streamableHttp.js";

import fs from "fs";

// 使用硅基流动的API配置
const client = new OpenAI({
  apiKey: "<替换为你的API Key>",  
  baseURL: "https://api.siliconflow.cn/v1"
});

/**
 * 单个MCP服务器连接管理类
 * 遵循单一职责原则，只负责单个服务器的连接和工具管理
 */
class MCPServerConnection {
  constructor(name, config) {
    this.name = name;
    this.config = config;
    this.mcp = new Client({ name: "mcp-client-with-config", version: "1.0.0" });
    this.transport = null;
    this.tools = [];
    this.connected = false;
  }

  async connect() {
    if (this.config.disabled) {
      console.log(`Server ${this.name} is disabled, skipping connection`);
      return;
    }

    try {
      // 远程MCP服务连接
      if (this.config.type === "sse") {
        this.transport = new StreamableHTTPClientTransport(
          new URL(this.config.url)
        );
      // 本地MCP服务连接
      }else if (this.config.type === "stdio") {
        this.transport = new StdioClientTransport({
          command: this.config.command,
          args: this.config.args,
        });
      }
      
      
      await this.mcp.connect(this.transport);
      
      const toolsResult = await this.mcp.listTools();
      this.tools = toolsResult.tools.map((tool) => ({
        name: tool.name,
        description: tool.description,
        input_schema: tool.inputSchema,
        serverName: this.name, // 添加服务器名称标识
      }));
      
      this.connected = true;
      console.log(`✓ 已连接到${this.config.type === "stdio"?"本地":"远程"}MCP服务器: ${this.name} 发现如下可用工具:`, this.tools.map(({ name }) => name));
    } catch (error) {
      console.error(`✗ Failed to connect to ${this.name}:`, error.message);
      throw error;
    }
  }

  async callTool(toolName, toolArguments) {
    if (!this.connected) {
      throw new Error(`Server ${this.name} is not connected`);
    }
    return await this.mcp.callTool({ name: toolName, arguments: toolArguments });
  }

  async disconnect() {
    if (this.transport) {
      await this.transport.close();
      this.connected = false;
    }
  }
}

/**
 * 多MCP服务器客户端管理类
 * 遵循开闭原则，可扩展支持更多服务器类型
 */
class MCPClient {
  constructor() {
    this.servers = new Map();
    this.allTools = [];
  }

  /**
   * 加载配置文件
   * @param {string} configPath - 配置文件路径
   * @returns {Object} 配置对象
   */
  loadConfig(configPath = "./mcp_settings.json") {
    try {
      const configData = fs.readFileSync(configPath, "utf8");
      return JSON.parse(configData);
    } catch (error) {
      console.error(`Failed to load config from ${configPath}:`, error.message);
      throw error;
    }
  }

  // 请求大模型
  async fetchLLM(messages){
    return await client.chat.completions.create({
      model: "Qwen/Qwen3-8B",
      max_tokens: 1000,
      enable_thinking: false,
      messages,
      tools: this.allTools.map(tool => ({
        type: 'function',
        function: {
          name: tool.name,
          description: tool.description,
          parameters: tool.input_schema
        }
      })),
    });
  }

  /**
   * 根据配置文件连接到多个MCP服务器
   * @param {string} configPath - 配置文件路径
   */
  async connectToServers(configPath = "./mcp_settings.json") {
    const config = this.loadConfig(configPath);
    const mcpServers = config.mcpServers || {};
    
    console.log(`🔗正通过配置装载 ${Object.keys(mcpServers).length} 个服务器...`);
    
    const connectionPromises = [];
    
    for (const [serverName, serverConfig] of Object.entries(mcpServers)) {
      if (!serverConfig.disabled) {
        const server = new MCPServerConnection(serverName, serverConfig);
        this.servers.set(serverName, server);
        connectionPromises.push(
          server.connect().catch(error => ({ serverName, error }))
        );
      } else {
        console.log(`⚠ Server ${serverName} is disabled in config`);
      }
    }
    
    // 并行连接所有服务器
    const results = await Promise.allSettled(connectionPromises);
    
    // 收集所有工具和统计连接结果
    this.allTools = [];
    let successfulConnections = 0;
    let failedConnections = 0;
    const failedServers = [];
    
    for (const server of this.servers.values()) {
      if (server.connected) {
        this.allTools.push(...server.tools);
        successfulConnections++;
      } else {
        failedConnections++;
        failedServers.push(server.name);
      }
    }
    
    console.log(`\n=== 连接状态 ===`);
    console.log(`✅ 成功连接: ${successfulConnections} 个服务器`);
    console.log(`❌ 连接失败: ${failedConnections} 个服务器`);
    console.log(`🛠️ 可用工具: ${this.allTools.length} 个`);
    
    if (this.allTools.length > 0) {
      console.log(`\n可用工具:`);
      this.allTools.forEach(tool => {
        console.log(`  - ${tool.name} (${tool.serverName})`);
      });
    }
    
    if (failedConnections > 0) {
      console.log(`\n连接失败:`);
      failedServers.forEach(serverName => {
        console.log(`  - ${serverName}`);
      });
    }
    
    console.log(`================\n`);
  }

  /**
   * 根据工具名称找到对应的服务器
   * @param {string} toolName - 工具名称
   * @returns {MCPServerConnection|null} 服务器连接对象
   */
  findToolServer(toolName) {
    for (const server of this.servers.values()) {
      if (server.tools.some(tool => tool.name === toolName)) {
        return server;
      }
    }
    return null;
  }

  /**
   * 使用Claude和可用工具处理问题
   * @param {string} query - 用户的输入问题
   * @returns {string} 处理后的响应
   */
  async processQuery(query) {
    if (this.allTools.length === 0) {
      throw new Error("No tools available. Please ensure at least one MCP server is connected.");
    }

    const messages = [
      {
        role: "user",
        content: query,
      },
    ];

    // 首次调用大模型，根据“用户的提问内容”结合“可用工具函数列表信息”进行【意图识别】并【给出工具函数调用决策】）
    const response = await this.fetchLLM(messages);

    // 处理响应和工具调用
    const message = response.choices[0].message;
    const finalText = [];
    const toolResults = [];

    // 如果有文本内容
    if (message.content) {
      finalText.push(message.content);
    }

    // 如果有工具调用
    if (message.tool_calls) {
      for (const toolCall of message.tool_calls) {
        const toolName = toolCall.function.name;
        const toolInput = JSON.parse(toolCall.function.arguments);
        
        console.log(`🔧 调用工具: ${toolName}`);
        console.log(`   参数: ${JSON.stringify(toolInput)}`);
        
        try {
          const server = this.findToolServer(toolName);
          if (!server) {
            throw new Error(`Tool ${toolName} not found in any connected server`);
          }
          
          console.log(`   服务器: ${server.name}`);
          const result = await server.callTool(toolName, toolInput);
          
          const toolOutput = result.content[0]?.text || JSON.stringify(result.content);
          console.log(`   ✅ 结果: ${toolOutput}`);
          
          toolResults.push({
            tool_call_id: toolCall.id,
            role: "tool",
            content: toolOutput,
          });
        } catch (error) {
          console.log(`   ❌ 错误: ${error.message}`);
          toolResults.push({
            tool_call_id: toolCall.id,
            role: "tool",
            content: `Error: ${error.message}`,
          });
        }
      }
    }

    // 如果有工具调用，需要再次调用API获取最终响应
    if (toolResults.length > 0) {
      messages.push({
        role: "assistant",
        content: message.content,
        tool_calls: message.tool_calls,
      });
      
      // 添加工具调用结果
      messages.push(...toolResults);

      // 再次调用大模型，整合【用户原始问题】、【函数调用决策】及【函数调用结果】给出最终结果。
      const finalResponse = await this.fetchLLM(messages);

      return finalResponse.choices[0].message.content || "";
    }

    return finalText.join("\n");
  }

  /**
   * 断开所有服务器连接
   */
  async disconnect() {
    console.log(`\n🔌 断开连接中...`);
    const disconnectPromises = [];
    for (const server of this.servers.values()) {
      if (server.connected) {
        console.log(`   断开 ${server.name}`);
        disconnectPromises.push(server.disconnect());
      }
    }
    await Promise.allSettled(disconnectPromises);
    this.servers.clear();
    this.allTools = [];
    console.log(`✅ 全部断开完成`);
  }

  /**
   * 获取连接状态信息
   * @returns {Object} 状态信息
   */
  getStatus() {
    const connectedServers = Array.from(this.servers.values()).filter(s => s.connected);
    return {
      totalServers: this.servers.size,
      connectedServers: connectedServers.length,
      totalTools: this.allTools.length,
      servers: connectedServers.map(s => ({
        name: s.name,
        toolCount: s.tools.length,
        tools: s.tools.map(t => t.name)
      }))
    };
  }
}

// 主程序
async function main() {
  const mcpClient = new MCPClient();
  
  try {
    // 根据配置文件连接到多个MCP服务器
    await mcpClient.connectToServers("./client/mcp_settings.json");
    
    // 显示连接状态
    const status = mcpClient.getStatus();
    if (status.connectedServers === 0) {
      console.error("❌ No servers connected. Please check your configuration.");
      return;
    }
    
    // 测试问题
    const prompts = [
      "99减2等于多少? 今天几号？今天是农历几号？",
    ];
    
    for (const prompt of prompts) {
      console.log(`\n🔍️ 问题: ${prompt}`);
      const result = await mcpClient.processQuery(prompt);
      console.log(`\n🎯 最终结果:`);
      console.log(`${'-'.repeat(50)}`);
      console.log(result);
      console.log(`${'-'.repeat(50)}\n`);
    }
  } catch (error) {
    console.error("❌ Error:", error);
  } finally {
    await mcpClient.disconnect();
  }
}

// 运行主程序
main().catch(console.error);